from functools import lru_cache
import torch
import os
from sklearn.metrics import accuracy_score
from torch.optim.lr_scheduler import StepLR

from dataset import Audio2ParamsDataset
from network import Audio2ParamsNetwork


if __name__ == '__main__':

    cwd = '/media/dajun/000EFE5F000190F2/dev_spa/DMuse/202202b3'
    
    config = {
            'dataset_config': {
                'train_sample_dir': cwd,
                'test_sample_dir': cwd,
                'train_label_dir': cwd,
                'test_label_dir': cwd,
                'compress_rate': 30.,
                'n_spaces': 30,
            },
            'batch_size': 2,
            'shuffle': True,
        }
    
    dataset_class = Audio2ParamsDataset
    train_dataset = dataset_class(mode='train', **config['dataset_config'])
    test_dataset = dataset_class(mode='test', **config['dataset_config'])
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config['batch_size'],
                                                 shuffle=config['shuffle'] if 'shuffle' in config else True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=config['batch_size'],
                                                shuffle=config['shuffle'] if 'shuffle' in config else True)
    
    # validate_func = validate_params
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if os.path.exists('DMuse/artear2_1/202202a_model.pth'):
        network = torch.load('DMuse/artear2_1/202202a_model.pth')
    else:
        network = Audio2ParamsNetwork(train_dataset.n_params, config['dataset_config']['n_spaces'])
    loss_fn = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(lr=0.0005, params=network.parameters())
    scheduler = StepLR(optimizer, step_size=1, gamma=0.5, verbose=True)
    epochs = 5
    validate_every = 1
    save_every = 20

    for epoch in range(epochs):
        for batch_id, (x_data, y_data) in enumerate(train_loader):
            # output = network(x_data.to(device)).to(torch.device('cpu'))
            output = network(x_data)

            pre_lab = torch.argmax(output, -1)
            # print(output.size(), pre_lab.size(), y_data.size())
            train_corrects = torch.sum(pre_lab == torch.argmax(y_data, -1).data) / train_dataset.n_params / config['batch_size']

            loss = loss_fn(output, y_data)
            optimizer.zero_grad()
            loss.backward()
            if batch_id % validate_every == 0:
                print("epoch: {}, batch_id: {}, loss: {}, acc: {}".format(epoch, batch_id, loss,
                                                                          train_corrects))
                # validate_func(network(next(test_loader())[1][0]), config['validation_config'])
                # torch.save(network.state_dict(), '202202a_model.pth')
            if batch_id % save_every == 0:
                torch.save(network, 'DMuse/artear2_1/202202a_model.pth')
            optimizer.step()
            scheduler.step()
